package com.artisan.day01;

import com.artisan.common.ApiKeys;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.mapdb.DB;
import org.mapdb.DBMaker;

import java.util.List;
import java.util.Map;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;
import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;
import static org.mapdb.Serializer.STRING;

/**
 * With AiServices:
 * - Persistent chat memory
 */
public class ServiceWithPersistentMemoryExample {


    interface Assistant {

        String chat(String message);
    }

    /**
     * 程序的入口点，用于演示如何使用AI助手进行聊天
     */
    public static void main(String[] args) {
        // 创建一个聊天记忆对象，设置最大消息数为10，并使用持久化聊天记忆存储
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .maxMessages(100)
                .chatMemoryStore(new PersistentChatMemoryStore())
                .build();

        // 创建一个聊天语言模型对象，设置API密钥和模型名称为GPT-4
        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(ApiKeys.OPENAI_API_KEY)
                .modelName(GPT_4_O_MINI)
                .build();

        // 使用AI服务创建一个助手对象，设置聊天语言模型和聊天记忆
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(model)
                .chatMemory(chatMemory)
                .build();

        // 发送消息给助手并打印回复
//        String answer = assistant.chat("你好，我的名字叫小工匠");
//        System.out.println(answer); // 你好，小工匠！很高兴认识你！有什么想聊的或者需要帮助的地方吗？


        // 现在，注释掉上面的两行代码，取消注释下面的两行代码，然后再次运行。
        // 这样做是为了演示助手如何根据之前的聊天记录记住用户的名字。

         String answerWithName = assistant.chat("我叫什么名字?");
         System.out.println(answerWithName); // 你叫小工匠！如果你有其他问题或者需要讨论的事情，随时告诉我！
    }

    /**
     * 持久化聊天内存存储类
     * 实现了ChatMemoryStore接口，用于持久化存储聊天消息
     * 该类使用MapDB库创建一个持久化的数据库，将聊天消息以键值对的形式存储在文件中
     */
    static class PersistentChatMemoryStore implements ChatMemoryStore {

        // 数据库实例，使用MapDB创建，启用事务处理
        private final DB db = DBMaker.fileDB("chat-memory.db").transactionEnable().make();

        // 存储聊天消息的Map，使用HashMap实现，键和值都是字符串类型
        private final Map<String, String> map = db.hashMap("messages", STRING, STRING).createOrOpen();

        /**
         * 根据memoryId获取聊天消息
         *
         * @param memoryId 聊天消息的唯一标识符
         * @return 聊天消息列表，如果找不到则返回null
         */
        @Override
        public List<ChatMessage> getMessages(Object memoryId) {
            String json = map.get((String) memoryId);
            return messagesFromJson(json);
        }

        /**
         * 更新特定memoryId的聊天消息列表
         *
         * @param memoryId 聊天消息的唯一标识符
         * @param messages 要更新的聊天消息列表
         */
        @Override
        public void updateMessages(Object memoryId, List<ChatMessage> messages) {
            String json = messagesToJson(messages);
            map.put((String) memoryId, json);
            db.commit();
        }

        /**
         * 删除特定memoryId的聊天消息
         *
         * @param memoryId 要删除的聊天消息的唯一标识符
         */
        @Override
        public void deleteMessages(Object memoryId) {
            map.remove((String) memoryId);
            db.commit();
        }
    }
}